import javax.swing.*;
import javax.swing.tree.DefaultMutableTreeNode;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class ID3_tree {
    private ArrayList<String> attribute = new ArrayList<>(); // 存储属性的名称
    private ArrayList<ArrayList<String>> attributevalue = new ArrayList<>(); // 存储每个属性的取值
    private ArrayList<String[]> data = new ArrayList<>(1730);// 原始数据
    int decatt; // 决策变量在属性集中的索引
    public static final String patternString = "@attribute(.*)[{](.*?)[}]";


    JTree jTree;
    DefaultMutableTreeNode root;
    File file;
    String relation = "relation";

    public ID3_tree(File file) {
        this.file = file;
    }

    JTree getJTree() {

        readARFF(file);
        root = new DefaultMutableTreeNode(relation);
        setDec("class");
        LinkedList<Integer> ll = new LinkedList<>();
        for (int i = 0; i < attribute.size(); i++) {
            if (i != decatt)
                ll.add(i);
        }
        ArrayList<Integer> al = new ArrayList<>();
        for (int i = 0; i < data.size(); i++) {
            al.add(i);
        }
        buildDT(root, al, ll);
        jTree = new JTree(root);
        return jTree;
    }


    //读取arff文件，给attribute、attributevalue、data赋值
    public void readARFF(File file) {
        try {
            FileReader fr = new FileReader(file);
            BufferedReader br = new BufferedReader(fr);
            String line;
            Pattern pattern = Pattern.compile(patternString);
            while ((line = br.readLine()) != null) {
                Matcher matcher = pattern.matcher(line);
                if (matcher.find()) {
                    attribute.add(matcher.group(1).trim());
                    String[] values = matcher.group(2).split(",");
                    ArrayList<String> al = new ArrayList<String>(values.length);
                    for (String value : values) {
                        al.add(value.trim());
                    }
                    attributevalue.add(al);
                } else if (line.startsWith("@data")) {
                    while ((line = br.readLine()) != null) {
                        if (line.equals(""))
                            continue;
                        String[] row = line.split(",");
                        data.add(row);
                    }
                } else {
                    if (line.startsWith("@relation"))
                        relation = line.substring(9);
                }
            }
            br.close();
        } catch (IOException e1) {
            e1.printStackTrace();
        }
    }

    //设置决策变量
    public void setDec(int n) {
        if (n < 0 || n >= attribute.size()) {
            JOptionPane.showMessageDialog(null,
                    "决策变量指定错误！\n请尝试把决策属性名改为“class”",
                    "error", JOptionPane.WARNING_MESSAGE);
            System.exit(2);
        }
        decatt = n;
    }

    public void setDec(String name) {
        int n = attribute.indexOf(name);
        setDec(n);
    }

    //给一个样本（数组中是各种情况的计数），计算它的熵
    public double getEntropy(int[] arr) {
        double entropy = 0.0;
        int sum = 0;
        for (int i = 0; i < arr.length; i++) {
            entropy -= arr[i] * Math.log(arr[i] + Double.MIN_VALUE) / Math.log(2);
            sum += arr[i];
        }
        entropy += sum * Math.log(sum + Double.MIN_VALUE) / Math.log(2);
        entropy /= sum;
        return entropy;
    }

    //给一个样本数组及样本的算术和，计算它的熵
    public double getEntropy(int[] arr, int sum) {
        double entropy = 0.0;
        for (int i = 0; i < arr.length; i++) {
            double u = 1.0 * arr[i] / sum;
            entropy -= u * Math.log(u + Double.MIN_VALUE) / Math.log(2);
        }
        return entropy;
    }

    public boolean infoPure(ArrayList<Integer> subset) {
        String value = data.get(subset.get(0))[decatt];
        for (int i = 1; i < subset.size(); i++) {
            String next = data.get(subset.get(i))[decatt];
            //equals表示对象内容相同，==表示两个对象指向的是同一片内存
            if (!value.equals(next))
                return false;
        }
        return true;
    }

    // 给定原始数据的子集(subset中存储行号),当以第index个属性为节点时计算它的信息熵
    public double calNodeEntropy(ArrayList<Integer> subset, int index) {
        int sum = subset.size();
        double entropy = 0.0;
        ArrayList<String> indexValue = attributevalue.get(index);//某属性
        int indexValueCount = indexValue.size();//某属性的取值个数
        int classCount = attributevalue.get(decatt).size();//分类结果的取值个数
        int[][] info = new int[indexValueCount][classCount];//某属性结果取值计数
        int[] count = new int[indexValueCount];//某属性计数
        //遍历子数据集
        for (Integer aSubset : subset) {
            int n = aSubset;
            String nodevalue = data.get(n)[index];//数据集中第n行该属性的取值
            int nodeind = indexValue.indexOf(nodevalue);
            count[nodeind]++;//该属性对应取值计数
            String decvalue = data.get(n)[decatt];//数据集中第n行分类结果取值
            int decind = attributevalue.get(decatt).indexOf(decvalue);
            info[nodeind][decind]++;//属性分类结果计数
        }
        for (int i = 0; i < info.length; i++) {
            entropy += getEntropy(info[i], count[i]) * count[i] / sum;
        }
        return entropy;
    }


    //
// 构建决策树
    public void buildDT(DefaultMutableTreeNode node, ArrayList<Integer> subset,
                        LinkedList<Integer> selatt) {


        if (infoPure(subset)) {
            node.setUserObject(node.getUserObject() + "," + data.get(subset.get(0))[decatt]);
            return;
        }
        int minIndex = -1;
        double minEntropy = Double.MAX_VALUE;
        //遍历属性集，找到最小的信息熵
        for (int i = 0; i < selatt.size(); i++) {
            if (i == decatt)
                continue;
            double entropy = calNodeEntropy(subset, selatt.get(i));
            if (entropy < minEntropy) {
                minIndex = selatt.get(i);
                minEntropy = entropy;
            }
        }
        String nodeName = attribute.get(minIndex);
        selatt.remove(new Integer(minIndex));

        ArrayList<String> attvalues = attributevalue.get(minIndex);//找当前决策的属性集
        for (String val : attvalues) {
            DefaultMutableTreeNode childNode = new DefaultMutableTreeNode("(" + nodeName + ":." + val + ")");
            node.add(childNode);
            //构造与当前属性取值相同的数据集
            ArrayList<Integer> al = new ArrayList<>();
            //从子数据集寻找新的数据集
            for (Integer aSubset : subset) {
                if (data.get(aSubset)[minIndex].equals(val)) {
                    al.add(aSubset);
                }
            }
            buildDT(childNode, al, selatt);
        }
        selatt.add(minIndex);//还原属性集
    }

}